Matrix block sum [Range Sum Query 2D - Immutable]

Time: O(MxN); Space: O(MxN); medium

Given a m * n matrix mat and an integer K, return a matrix answer where each answer[i][j] is the sum of all elements mat[r][c] for i - K <= r <= i + K, j - K <= c <= j + K, and (r, c) is a valid position in the matrix.

Example 1:

Input: mat = [[1,2,3],[4,5,6],[7,8,9]], K = 1

Output: [[12,21,16],[27,45,33],[24,39,28]]

Example 2:

Input: mat = [[1,2,3],[4,5,6],[7,8,9]], K = 2

Output: [[45,45,45],[45,45,45],[45,45,45]]

Constraints:

  • m = len(mat)

  • n = len(mat[i])

  • 1 <= m, n, K <= 100

  • 1 <= mat[i][j] <= 100

Hints:

  1. How to calculate the required sum for a cell (i,j) fast ?

  2. Use the concept of cumulative sum array.

  3. Create a cumulative sum matrix where dp[i][j] is the sum of all cells in the rectangle from (0,0) to (i,j), use inclusion-exclusion idea.

1. Dynamic programming (Range Sum Query) [O(MxN), O(MxN)]

[1]:
class Solution1(object):
    """
    Time: O(M*N)
    Space: O(M*N)
    """
    def matrixBlockSum(self, mat, K):
        """
        :type mat: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        m, n = len(mat), len(mat[0])
        accu = [[0 for _ in range(n+1)] for _ in range(m+1)]

        for i in range(m):
            for j in range(n):
                accu[i+1][j+1] = accu[i+1][j]+accu[i][j+1]-accu[i][j]+mat[i][j]

        result = [[0 for _ in range(n)] for _ in range(m)]

        for i in range(m):
            for j in range(n):
                r1, c1, r2, c2 = max(i-K, 0), max(j-K, 0), min(i+K+1, m), min(j+K+1, n)
                result[i][j] = accu[r2][c2]-accu[r1][c2]-accu[r2][c1]+accu[r1][c1]

        return result
[2]:
s = Solution1()

mat = [[1,2,3],[4,5,6],[7,8,9]]
K = 1
assert s.matrixBlockSum(mat, K) == [[12,21,16],[27,45,33],[24,39,28]]

mat = [[1,2,3],[4,5,6],[7,8,9]]
K = 2
assert s.matrixBlockSum(mat, K) == [[45,45,45],[45,45,45],[45,45,45]]